/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.iotdb.confignode.it.cluster;

import org.apache.iotdb.common.rpc.thrift.*;
import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
import org.apache.iotdb.confignode.rpc.thrift.*;
import org.apache.iotdb.consensus.ConsensusFactory;
import org.apache.iotdb.it.env.ConfigFactory;
import org.apache.iotdb.it.env.ConfigNodeWrapper;
import org.apache.iotdb.it.env.DataNodeWrapper;
import org.apache.iotdb.it.env.EnvFactory;
import org.apache.iotdb.it.framework.IoTDBTestRunner;
import org.apache.iotdb.itbase.category.ClusterIT;
import org.apache.iotdb.itbase.env.BaseConfig;
import org.apache.iotdb.rpc.TSStatusCode;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static org.junit.Assert.*;

@RunWith(IoTDBTestRunner.class)
@Category({ClusterIT.class})
public class IoTDBClusterNodeGetterIT {

    private static final BaseConfig CONF = ConfigFactory.getConfig();

    private static final int testConfigNodeNum = 2;
    private static final int testDataNodeNum = 2;

    protected static String originalConfigNodeConsensusProtocolClass;
    protected static String originalSchemaRegionConsensusProtocolClass;
    protected static String originalDataRegionConsensusProtocolClass;
    private static final String testConsensusProtocolClass = ConsensusFactory.RATIS_CONSENSUS;

    @Before
    public void setUp() throws Exception {
        originalConfigNodeConsensusProtocolClass = CONF.getConfigNodeConsensusProtocolClass();
        CONF.setConfigNodeConsesusProtocolClass(testConsensusProtocolClass);
        originalSchemaRegionConsensusProtocolClass = CONF.getSchemaRegionConsensusProtocolClass();
        CONF.setSchemaRegionConsensusProtocolClass(testConsensusProtocolClass);
        originalDataRegionConsensusProtocolClass = CONF.getDataRegionConsensusProtocolClass();
        CONF.setDataRegionConsensusProtocolClass(testConsensusProtocolClass);

        // Init 2C2D environment
        EnvFactory.getEnv().initClusterEnvironment(testConfigNodeNum, testDataNodeNum);
    }

    @After
    public void tearDown() {
        EnvFactory.getEnv().cleanAfterClass();
        CONF.setConfigNodeConsesusProtocolClass(originalConfigNodeConsensusProtocolClass);
        CONF.setSchemaRegionConsensusProtocolClass(originalSchemaRegionConsensusProtocolClass);
        CONF.setDataRegionConsensusProtocolClass(originalDataRegionConsensusProtocolClass);
    }

    @Test
    public void showClusterAndNodesTest() throws Exception {
        try (SyncConfigNodeIServiceClient client =
                     (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {

            List<ConfigNodeWrapper> configNodeWrappers = EnvFactory.getEnv().getConfigNodeWrapperList();
            List<DataNodeWrapper> dataNodeWrappers = EnvFactory.getEnv().getDataNodeWrapperList();

            /* Test showCluster */
            TShowClusterResp showClusterResp = client.showCluster();
            // Check ConfigNodeLocation
            List<TConfigNodeLocation> configNodeLocations = showClusterResp.getConfigNodeList();
            for (TConfigNodeLocation configNodeLocation : configNodeLocations) {
                boolean found = false;
                for (ConfigNodeWrapper configNodeWrapper : configNodeWrappers) {
                    if (configNodeWrapper.getIp().equals(configNodeLocation.getInternalEndPoint().getIp())
                            && configNodeWrapper.getPort() == configNodeLocation.getInternalEndPoint().getPort()
                            && configNodeWrapper.getConsensusPort()
                            == configNodeLocation.getConsensusEndPoint().getPort()) {
                        found = true;
                        break;
                    }
                }
                assertTrue(found);
            }
            // Check DataNodeLocation
            List<TDataNodeLocation> dataNodeLocations = showClusterResp.getDataNodeList();
            for (TDataNodeLocation dataNodeLocation : dataNodeLocations) {
                boolean found = false;
                for (DataNodeWrapper dataNodeWrapper : dataNodeWrappers) {
                    if (dataNodeWrapper.getIp().equals(dataNodeLocation.getInternalEndPoint().getIp())
                            && dataNodeWrapper.getPort() == dataNodeLocation.getClientRpcEndPoint().getPort()
                            && dataNodeWrapper.getInternalPort()
                            == dataNodeLocation.getInternalEndPoint().getPort()
                            && dataNodeWrapper.getMppDataExchangePort()
                            == dataNodeLocation.getMPPDataExchangeEndPoint().getPort()
                            && dataNodeWrapper.getSchemaRegionConsensusPort()
                            == dataNodeLocation.getSchemaRegionConsensusEndPoint().getPort()
                            && dataNodeWrapper.getDataRegionConsensusPort()
                            == dataNodeLocation.getDataRegionConsensusEndPoint().getPort()) {
                        found = true;
                        break;
                    }
                }
                assertTrue(found);
            }

            /* Tests showClusterParameters */
            TShowClusterParametersResp showClusterParametersResp = client.showClusterParameters();
            Assert.assertEquals(
                    TSStatusCode.SUCCESS_STATUS.getStatusCode(),
                    showClusterParametersResp.getStatus().getCode());
            TClusterParameters clusterParameters = showClusterParametersResp.getClusterParameters();
            Assert.assertEquals(
                    testConsensusProtocolClass, clusterParameters.getConfigNodeConsensusProtocolClass());
            Assert.assertEquals(
                    testConsensusProtocolClass, clusterParameters.getDataRegionConsensusProtocolClass());
            Assert.assertEquals(
                    testConsensusProtocolClass, clusterParameters.getSchemaRegionConsensusProtocolClass());
            Assert.assertEquals(
                    CONF.getSeriesPartitionSlotNum(), clusterParameters.getSeriesPartitionSlotNum());
            Assert.assertEquals(
                    CONF.getSeriesPartitionExecutorClass(),
                    clusterParameters.getSeriesPartitionExecutorClass());
            Assert.assertEquals(CONF.getDefaultTTL(), clusterParameters.getDefaultTTL());
            Assert.assertEquals(
                    CONF.getTimePartitionInterval(), clusterParameters.getTimePartitionInterval());
            Assert.assertEquals(
                    CONF.getDataReplicationFactor(), clusterParameters.getDataReplicationFactor());
            Assert.assertEquals(
                    CONF.getSchemaReplicationFactor(), clusterParameters.getSchemaReplicationFactor());
            Assert.assertEquals(
                    CONF.getDataRegionPerProcessor(), clusterParameters.getDataRegionPerProcessor(), 0.01);
            Assert.assertEquals(
                    CONF.getSchemaRegionPerDataNode(), clusterParameters.getSchemaRegionPerDataNode(), 0.01);
            Assert.assertEquals(
                    CONF.getDiskSpaceWarningThreshold(),
                    clusterParameters.getDiskSpaceWarningThreshold(),
                    0.01);
            Assert.assertEquals(
                    CONF.getReadConsistencyLevel(), clusterParameters.getReadConsistencyLevel());
            Assert.assertEquals(
                    CONF.getLeastDataRegionGroupNum(), clusterParameters.getLeastDataRegionGroupNum());

            /* Test showConfigNodes */
            TShowConfigNodesResp showConfigNodesResp = client.showConfigNodes();
            // Check ConfigNodeInfo
            List<TConfigNodeInfo> configNodesInfo = showConfigNodesResp.getConfigNodesInfoList();
            for (TConfigNodeInfo configNodeInfo : configNodesInfo) {
                boolean found = false;
                for (ConfigNodeWrapper configNodeWrapper : configNodeWrappers) {
                    if (configNodeWrapper.getIp().equals(configNodeInfo.getInternalAddress())
                            && configNodeWrapper.getPort() == configNodeInfo.getInternalPort()) {
                        found = true;
                        break;
                    }
                }
                assertTrue(found);
            }

            /* Test showDataNodes */
            TShowDataNodesResp showDataNodesResp = client.showDataNodes();
            // Check DataNodeInfo
            List<TDataNodeInfo> dataNodesInfo = showDataNodesResp.getDataNodesInfoList();
            for (TDataNodeInfo dataNodeInfo : dataNodesInfo) {
                boolean found = false;
                for (DataNodeWrapper dataNodeWrapper : dataNodeWrappers) {
                    if (dataNodeWrapper.getIp().equals(dataNodeInfo.getRpcAddresss())
                            && dataNodeWrapper.getPort() == dataNodeInfo.getRpcPort()) {
                        found = true;
                        break;
                    }
                }
                assertTrue(found);
            }
        }
    }

    @Test
    public void removeAndStopConfigNodeTest() throws Exception {
        TShowClusterResp showClusterResp;
        TSStatus status;

        try (SyncConfigNodeIServiceClient client =
                     (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {

            // Test remove ConfigNode
            showClusterResp = client.showCluster();
            TConfigNodeLocation removedConfigNodeLocation = showClusterResp.getConfigNodeList().get(1);
            status = client.removeConfigNode(removedConfigNodeLocation);
            Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());

            // Waiting for RemoveConfigNodeProcedure
            boolean isRemoved = false;
            for (int retry = 0; retry < 10; retry++) {
                showClusterResp = client.showCluster();
                if (showClusterResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
                        && showClusterResp.getConfigNodeListSize() == testConfigNodeNum - 1) {
                    isRemoved = true;
                    break;
                }

                TimeUnit.SECONDS.sleep(1);
            }
            if (!isRemoved) {
                fail("Remove ConfigNode failed");
            }

            // Test stop ConfigNode
            status = client.stopConfigNode(removedConfigNodeLocation);
            assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
        }
    }

    @Test
    public void queryAndRemoveDataNodeTest() throws Exception {

        try (SyncConfigNodeIServiceClient client =
                     (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) {

            // Pick a registered DataNode
            TDataNodeLocation dataNodeLocation = new TDataNodeLocation();
            TShowDataNodesResp showDataNodesResp = client.showDataNodes();
            TDataNodeInfo dataNodeInfo = showDataNodesResp.getDataNodesInfoList().get(0);

            dataNodeLocation.setDataNodeId(dataNodeInfo.getDataNodeId());
            dataNodeLocation.setClientRpcEndPoint(
                    new TEndPoint(dataNodeInfo.getRpcAddresss(), dataNodeInfo.getRpcPort()));
            dataNodeLocation.setMPPDataExchangeEndPoint(
                    new TEndPoint(dataNodeInfo.getRpcAddresss(), dataNodeInfo.getRpcPort() + 1));
            dataNodeLocation.setInternalEndPoint(
                    new TEndPoint(dataNodeInfo.getRpcAddresss(), dataNodeInfo.getRpcPort() + 2));
            dataNodeLocation.setDataRegionConsensusEndPoint(
                    new TEndPoint(dataNodeInfo.getRpcAddresss(), dataNodeInfo.getRpcPort() + 3));
            dataNodeLocation.setSchemaRegionConsensusEndPoint(
                    new TEndPoint(dataNodeInfo.getRpcAddresss(), dataNodeInfo.getRpcPort() + 4));

            /* Test query one DataNodeInfo */
            TDataNodeConfigurationResp dataNodeConfigurationResp =
                    client.getDataNodeConfiguration(dataNodeLocation.getDataNodeId());
            Map<Integer, TDataNodeConfiguration> configurationMap =
                    dataNodeConfigurationResp.getDataNodeConfigurationMap();
            TDataNodeLocation dnLocation =
                    dataNodeConfigurationResp
                            .getDataNodeConfigurationMap()
                            .get(dataNodeLocation.getDataNodeId())
                            .getLocation();
            assertEquals(
                    TSStatusCode.SUCCESS_STATUS.getStatusCode(),
                    dataNodeConfigurationResp.getStatus().getCode());
            assertEquals(1, configurationMap.size());
            assertEquals(dataNodeLocation, dnLocation);

            /* Test query all DataNodeConfiguration */
            dataNodeConfigurationResp = client.getDataNodeConfiguration(-1);
            configurationMap = dataNodeConfigurationResp.getDataNodeConfigurationMap();
            assertEquals(
                    TSStatusCode.SUCCESS_STATUS.getStatusCode(),
                    dataNodeConfigurationResp.getStatus().getCode());
            assertEquals(testDataNodeNum, configurationMap.size());

            /* Test remove DataNode */
            // Remove DataNode
            dataNodeLocation =
                    dataNodeConfigurationResp
                            .getDataNodeConfigurationMap()
                            .get(
                                    showDataNodesResp.getDataNodesInfoList().get(testDataNodeNum - 1).getDataNodeId())
                            .getLocation();
            TDataNodeRemoveReq dataNodeRemoveReq =
                    new TDataNodeRemoveReq(Collections.singletonList(dataNodeLocation));
            TDataNodeRemoveResp dataNodeRemoveResp = client.removeDataNode(dataNodeRemoveReq);
            assertEquals(
                    TSStatusCode.SUCCESS_STATUS.getStatusCode(), dataNodeRemoveResp.getStatus().getCode());

            // Waiting for RemoveDataNodeProcedure
            for (int retry = 0; retry < 10; retry++) {
                showDataNodesResp = client.showDataNodes();
                if (showDataNodesResp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
                        && showDataNodesResp.getDataNodesInfoListSize() == testDataNodeNum - 1) {
                    return;
                }

                TimeUnit.SECONDS.sleep(1);
            }
            fail("Remove DataNode failed");
        }
    }
}
